from sklearn.feature_selection import SelectKBest, chi2
from sklearn.base import BaseEstimator, TransformerMixin
import numpy
__author__ = 'panagiotis'


class UniFeatureSelector(BaseEstimator, TransformerMixin):

    def __init__(self, max_features=500):
        self.n_features = max_features
        self.features = set()
        self.mask = None

    def fit(self, X, y):
        self.mask = numpy.zeros((X.shape[1],), dtype=bool)
        for cl in set(y):
            feat_selection = SelectKBest(chi2, k=self.n_features / len(set(y)))
            feat_selection.fit(X, [1 if r == cl else 0 for r in y])
            self.features |= set(feat_selection.get_support(indices=True))
        return self

    def get_features(self, indices=False):
        if indices:
            return sorted(self.features)
        else:
            self.mask[sorted(self.features)] = True
            return self.mask

    def transform(self, X):
        return X[:, sorted(self.features)]

    def inverse_transform(self, X):
        Z = numpy.zeros((X.shape[0], self.mask.shape[0]))
        Z[:, self.get_features(indices=True)] = X
        return Z
